// ============================================================================
// File: pkg/metrics/metrics.go - Prometheus Metrics
// ============================================================================
package metrics

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"google.golang.org/grpc"
	"context"
)

var (
	// Snapshot operations
	SnapshotOperations = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "onyx_snapshot_operations_total",
			Help: "Total number of snapshot operations",
		},
		[]string{"operation"},
	)

	SnapshotDuration = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "onyx_snapshot_duration_seconds",
			Help:    "Duration of snapshot operations",
			Buckets: prometheus.DefBuckets,
		},
		[]string{"operation"},
	)

	SnapshotErrors = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "onyx_snapshot_errors_total",
			Help: "Total number of snapshot errors",
		},
		[]string{"operation"},
	)

	ActiveSnapshots = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "onyx_active_snapshots",
			Help: "Number of active snapshots",
		},
	)

	// Deduplication metrics
	DedupSavedBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_dedup_saved_bytes_total",
			Help: "Total bytes saved by deduplication",
		},
	)

	DedupHitRate = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "onyx_dedup_hits_total",
			Help: "Total deduplication hits",
		},
		[]string{"level"}, // file or block
	)

	DedupMissRate = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "onyx_dedup_misses_total",
			Help: "Total deduplication misses",
		},
		[]string{"level"},
	)

	StoredBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_stored_bytes_total",
			Help: "Total bytes stored (after dedup)",
		},
	)

	// Cache metrics
	CacheHits = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_cache_hits_total",
			Help: "Total cache hits",
		},
	)

	CacheMisses = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_cache_misses_total",
			Help: "Total cache misses",
		},
	)

	CacheSize = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "onyx_cache_size_bytes",
			Help: "Current cache size in bytes",
		},
	)

	CacheEvictions = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_cache_evictions_total",
			Help: "Total cache evictions",
		},
	)

	// Pull metrics
	LayerPulls = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_layer_pulls_total",
			Help: "Total layer pulls",
		},
	)

	PullDuration = promauto.NewHistogram(
		prometheus.HistogramOpts{
			Name:    "onyx_pull_duration_seconds",
			Help:    "Duration of layer pulls",
			Buckets: prometheus.DefBuckets,
		},
	)

	PullErrors = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_pull_errors_total",
			Help: "Total pull errors",
		},
	)

	DownloadedBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_downloaded_bytes_total",
			Help: "Total bytes downloaded",
		},
	)

	// Prefetch metrics
	PrefetchJobs = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_prefetch_jobs_total",
			Help: "Total prefetch jobs",
		},
	)

	PrefetchDuration = promauto.NewHistogram(
		prometheus.HistogramOpts{
			Name:    "onyx_prefetch_duration_seconds",
			Help:    "Duration of prefetch operations",
			Buckets: prometheus.DefBuckets,
		},
	)

	PrefetchedFiles = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_prefetched_files_total",
			Help: "Total files prefetched",
		},
	)

	PrefetchedBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "onyx_prefetched_bytes_total",
			Help: "Total bytes prefetched",
		},
	)
)

type Timer struct {
	start time.Time
}

func NewTimer() *Timer {
	return &Timer{start: time.Now()}
}

func (t *Timer) ObserveDuration() time.Duration {
	return time.Since(t.start)
}

func Init() {
	// Register all metrics
	prometheus.MustRegister(
		SnapshotOperations,
		SnapshotDuration,
		SnapshotErrors,
		ActiveSnapshots,
		DedupSavedBytes,
		DedupHitRate,
		DedupMissRate,
		StoredBytes,
		CacheHits,
		CacheMisses,
		CacheSize,
		CacheEvictions,
		LayerPulls,
		PullDuration,
		PullErrors,
		DownloadedBytes,
		PrefetchJobs,
		PrefetchDuration,
		PrefetchedFiles,
		PrefetchedBytes,
	)
}

// UnaryServerInterceptor returns a gRPC interceptor for metrics
func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, 
		handler grpc.UnaryHandler) (interface{}, error) {
		
		timer := NewTimer()
		resp, err := handler(ctx, req)
		duration := timer.ObserveDuration()

		// Record metrics
		if err != nil {
			SnapshotErrors.WithLabelValues(info.FullMethod).Inc()
		}
		SnapshotDuration.WithLabelValues(info.FullMethod).Observe(duration.Seconds())

		return resp, err
	}
}